# coding=utf-8
# Copyright (C) yhy team - All Rights Reserved
#
# @Version:   3.10.4
# @Software:  PyCharm
# @FileName:  vitebi.py
# @CTime:     2022/10/12 15:13   
# @Author:    yhy
# @Email:     yhy@yhy.com
# @UTime:     2022/10/12 15:13
#
# @Description:
#     
#     xxx
#
import codecs
import logging
from typing import List, Dict, Optional
import numpy as np


logger = logging.getLogger(__name__)


def demo1():
    # Initial Probabilities
    p_s, p_r = 2 / 3, 1 / 3

    # Transition Probabilities
    p_ss, p_sr, p_rs, p_rr, = 0.8, 0.2, 0.4, 0.6

    # Emission Probabilities
    p_sh, p_sg, p_rh, p_rg = 0.8, 0.2, 0.4, 0.6

    # instances
    # instance1
    mood = ['H', 'H', 'G', 'G', 'G', 'H']
    proba = []
    weather = []

    if mood[0] == 'H':
        proba.append((p_s * p_sh, p_r * p_rh))
    else:
        proba.append((p_s * p_sg, p_r * p_rg))

    for i in range(1, len(mood)):
        sunny_1, rainy_1 = proba[-1]
        if mood[i] == 'H':
            sunny0 = max(sunny_1 * p_ss * p_sh, rainy_1 * p_rs * p_sh)
            rainy0 = max(sunny_1 * p_sr * p_rh, rainy_1 * p_rr * p_rh)
        else:
            sunny0 = max(sunny_1 * p_ss * p_sg, rainy_1 * p_rs * p_sg)
            rainy0 = max(sunny_1 * p_sr * p_rg, rainy_1 * p_rr * p_rg)
        proba.append((sunny0, rainy0))

    for p in proba:
        if p[0] > p[1]:
            weather.append('S')
        else:
            weather.append('R')

    return weather


def viterbi(states, observations, prior_probability, transition, emission, observation_seqs):
    n, m, t = len(states), len(observation_seqs), len(observation_seqs)
    probability = np.ones((t, n))

    obs = [observations.index(o) for o in observation_seqs]

    probability[0, :] = prior_probability * emission[:, obs[0]]

    for i, o in enumerate(obs[1:], 1):
        for j, s in enumerate(states):
            probability[i, j] = np.max(probability[i-1, :] * transition[:, j] * emission[j, o])

    state_seq = [states[np.argmax(p_vec)] for p_vec in probability]

    return probability, state_seq


def demo2():
    states = ['sunny', 'rainy']
    observations = ['Happy', 'Grumpy']
    prior_probability = [2/3, 1/3]
    transition = np.array([[0.8, 0.2], [0.4, 0.6]])
    emission = np.array([[0.8, 0.2], [0.4, 0.6]])
    observation_seqs = ['Happy', 'Happy', 'Grumpy', 'Grumpy', 'Grumpy', 'Happy']

    probability, state_seq = viterbi(states, observations, prior_probability, transition, emission, observation_seqs)

    return state_seq
if __name__ == '__main__':

    outs = demo2()
    print(outs)